from PIL import Image
from torchvision import models, transforms
import torch
from torchvision.models import ResNet50_Weights


# 预处理函数：调整大小、灰度化
def preprocess_image(image_path):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # 调整为模型输入尺寸
        transforms.Grayscale(num_output_channels=3),  # 转为3通道灰度图
        transforms.ToTensor(),  # 转为Tensor
        transforms.Normalize([0.5], [0.5])  # 归一化
    ])
    image = Image.open(image_path).convert("RGB")
    return transform(image).unsqueeze(0)  # 增加batch维度


# 使用预训练模型提取图片向量
def get_image_vector(tensor):
    # 使用新的 weights 参数加载预训练模型
    model = models.resnet50(weights=ResNet50_Weights.DEFAULT)
    model.fc = torch.nn.Identity()  # 移除全连接层
    model.eval()  # 设置为评估模式
    with torch.no_grad():
        vector = model(tensor)
    return vector.numpy()  # 返回为 Numpy 数组


if __name__ == "__main__":
    # 示例
    image_tensor = preprocess_image("../static/ai.png")
    print(image_tensor)
    vector1 = get_image_vector(image_tensor)
    print("Image Vector:", vector1)

    # 得到的 Numpy 数组可以将其转化为 json 存到数据库的字段中，也可以使用专业的向量数据库进行存储。
